import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
from PyCmpltrtok.common import sep

moive = ['Love Movie', 'Action Movie']

T = np.array([[3, 104, 0],
     [2, 100, 0],
     [1, 81, 0,],
     [101, 10, 1],
     [99, 5, 1],
     [98, 2, 1]], dtype=np.float64)
m = len(T)

x = np.array([18, 90])
k = 3

cmap = plt.cm.get_cmap('rainbow', 2)
idx_neg = T[:, 2] == 0
idx_pos = np.invert(idx_neg)
m_neg = np.sum(idx_neg)
m_pos = np.sum(idx_pos)
c00 = cmap(0)
c0 = np.tile(cmap(0), m_neg)
plt.scatter(T[idx_neg, 0], T[idx_neg, 1], c=np.atleast_2d(cmap(0)), label='negative')
plt.scatter(T[idx_pos, 0], T[idx_pos, 1], c=np.atleast_2d(cmap(1)), label='positive')
plt.scatter(x[0], x[1], c='b')
plt.legend()

dis = np.zeros([m, 2])
# for i, t in enumerate(T):
#     # 欧氏距离
#     d = np.sqrt((t[0] - x[0])**2 + (t[1] - x[1])**2)
#     dis[i, 0] = d
#     dis[i, 1] = t[2]
dx = T[:, 0] - x[0]
dy = T[:, 1] - x[1]
dx2 = dx**2
dy2 = dy**2
dsum = dx2 + dy2
sqrt = np.sqrt(dsum)
dis[:, 0] = sqrt
dis[:, 1] = T[:, 2]

sep('distances')
arg = np.argsort(dis[:, 0])
dis = dis[arg]
print(dis)
dis = dis[:k]

sep('Get category')
# This approach is too complex
# counts = np.bincount(dis[:, 1].astype(np.int64))
# print(counts)
# h= counts.argmax()

s = pd.Series(dis[:, 1])
# s = pd.Series([0, 1, 0, 1, 1, 1, 1])  # test
sc = s.value_counts()
print(sc)
h = sc.index[0]

print(f'Hypothesis = {h}')

plt.show()
